import torch
import torch.nn as nn
import torch.optim as optim


# 1. 定义神经网络架构
class CustomAgentModel(nn.Module):
    """自定义智能体决策模型"""

    def __init__(self, input_size, hidden_size, output_size):
        super(CustomAgentModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)  # 输入层
        self.relu = nn.ReLU()  # 激活函数
        self.fc2 = nn.Linear(hidden_size, output_size)  # 输出层

    def forward(self, x):
        """前向传播"""
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x


# 2. 创建模型实例
input_size = 10  # 输入特征维度
hidden_size = 32  # 隐藏层大小
output_size = 5  # 输出动作空间
model = CustomAgentModel(input_size, hidden_size, output_size)

# 3. 准备数据
inputs = torch.randn(100, input_size)  # 100个样本的模拟输入
labels = torch.randint(0, output_size, (100,))  # 随机标签

# 4. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器

# 5. 训练循环
for epoch in range(100):
    # 前向传播
    outputs = model(inputs)

    # 计算损失
    loss = criterion(outputs, labels)

    # 反向传播
    optimizer.zero_grad()  # 清空梯度
    loss.backward()  # 反向计算梯度
    optimizer.step()  # 更新权重

    # 打印训练进度
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/100], Loss: {loss.item():.4f}')

# 6. 模型推理
test_input = torch.randn(1, input_size)  # 新输入
predicted = model(test_input)
print(f"预测结果: {torch.argmax(predicted).item()}")